# services/qa_service.py
from services.retrieval_service import get_retriever
from services.llm_service import deepseek_llm
from prompts.financial_prompt import create_prompt
import re
import json

# 全局变量存储RAG处理链
rag_chain = None


def build_rag_chain():
    """构建RAG处理链"""
    retriever = get_retriever()

    def rag_pipeline(question):
        # 1. 检索相关上下文
        context_docs = retriever.get_relevant_documents(question)
        context = "\n\n".join([doc.page_content for doc in context_docs])

        # 2. 创建提示词
        prompt = create_prompt(context, question)

        # 3. 调用DeepSeek生成
        return deepseek_llm.generate_json(prompt)

    return rag_pipeline


def initialize_rag_chain():
    """初始化RAG处理链"""
    global rag_chain
    if rag_chain is None:
        rag_chain = build_rag_chain()


def is_rag_chain_ready():
    """检查RAG处理链是否就绪"""
    return rag_chain is not None


def ask_question_service(question: str):
    """问答服务逻辑"""
    global rag_chain

    if rag_chain is None:
        raise Exception("RAG处理链未初始化")

    # 实时数据处理（股票代码检测）
    stock_codes = re.findall(r"(sh\d{6}|sz\d{6})", question.lower())
    if stock_codes:
        # 注意：这里需要你实现 get_stock_data 函数
        # realtime_data = get_stock_data(stock_codes[0].upper())
        # question += f"\n[实时数据] {json.dumps(realtime_data, ensure_ascii=False)}"
        pass  # 暂时跳过实时数据处理，因为需要你实现 get_stock_data

    # 执行RAG流程
    response = rag_chain(question)
    return response
